from lib import *

def test_convolution():
    in_height = np.random.randint(1, 100)
    in_width = np.random.randint(1, 100)
    in_channel = np.random.randint(3, 100)
    out_channel = np.random.randint(3, 100)
    H_L1 = np.random.randint(20,100)
    W_L1 = np.random.randint(20,100)
    IC_L1 = ceil_div(np.random.randint(2, in_channel), K0)*K0
    OC_L1 = ceil_div(np.random.randint(2, out_channel), N0)*N0
    H_L0 = np.random.randint(10,20)
    W_L0 = np.random.randint(10,20)
    IC_L0 = K0*np.random.randint(1,5)
    OC_L0 = N0*np.random.randint(1,5)
    kernel_h = np.random.randint(1,5)
    kernel_w = np.random.randint(1,5)
    stride_h = np.random.randint(1,5)
    stride_w = np.random.randint(1,5)
    pad_h = np.random.randint(0, kernel_h)
    pad_w = np.random.randint(0, kernel_w)
    dilation_h = np.random.randint(1, in_height+1)
    dilation_w = np.random.randint(1, in_width+1)
    out_height = (in_height + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) // stride_h + 1
    out_width = (in_width + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) // stride_w + 1
    if(out_height <= 0 or out_width <= 0):
        return False
    # 定义卷积层参数
    layer = {
    'in_height': in_height,
    'in_width': in_width,
    'out_height': out_height,
    'out_width': out_width,
    'in_channel': in_channel,
    'out_channel': out_channel,
    'kernel_h': kernel_h,
    'kernel_w': kernel_w,
    'stride_h': stride_h,
    'stride_w': stride_w,
    'pad_h': pad_h,
    'pad_w': pad_w,
    'dilation_h': dilation_h,
    'dilation_w': dilation_w
    }
    print(layer)
    print("H_L1: ", H_L1, ", W_L1: ", W_L1, ", IC_L1: ", IC_L1, ", OC_L1: ", OC_L1, ", H_L0: ", H_L0, ", W_L0: ", W_L0, ", IC_L0: ", IC_L0, ", OC_L0: ", OC_L0)

    # 创建输入数据
    ifm = np.random.randint(-128, 127, size=(in_channel, in_height, in_width))
    weight = np.random.randint(-128, 127, size=(out_channel, in_channel, kernel_h, kernel_w))
    bias = np.random.randint(-128, 127, size=out_channel)
    deq = np.random.rand(out_channel)
    ofm = np.zeros((out_channel, out_height, out_width))
    weight_MK = np.zeros((out_channel, ceil_div(in_channel, K0)*K0, kernel_h, kernel_w))
    weight_MK[:, :in_channel, :, :] = weight
    weight_MK1K0 = weight_MK.reshape(out_channel, ceil_div(in_channel, K0), K0, kernel_h, kernel_w).transpose((0, 1, 3, 4, 2))
    weight_MK = weight_MK1K0.reshape(out_channel, ceil_div(in_channel, K0)*kernel_h*kernel_w*K0)

    # 先转成CHWc, 再使用im2col和gemm
    ifm_HWC = ifm.transpose((1, 2, 0))
    for l1_oc_start_in_l2 in range(0, out_channel, OC_L1):
        for l1_oh_start_in_l2 in range(0, out_height, H_L1):
            for l1_ow_start_in_l2 in range(0, out_width, W_L1):
                oc_size_l1 = min(OC_L1, out_channel - l1_oc_start_in_l2)
                oh_size_l1 = min(H_L1, out_height - l1_oh_start_in_l2)
                ow_size_l1 = min(W_L1, out_width - l1_ow_start_in_l2)
                M = H_L0 * W_L0
                N = OC_L0
                M1 = ceil_div(M, M0)
                N1 = ceil_div(N, N0)
                result_m2n2m1n1m0n0_psb = np.zeros((ceil_div(oh_size_l1, H_L0)*ceil_div(ow_size_l1, W_L0)*ceil_div(oc_size_l1, OC_L0), M1, N1, M0, N0))
                for l1_ic_start_in_l2 in range(0, in_channel, IC_L1):
                    l1_oh_end_in_l2 = min(l1_oh_start_in_l2 + H_L1, out_height) - 1
                    l1_ow_end_in_l2 = min(l1_ow_start_in_l2 + W_L1, out_width) - 1
                    l1_ih_start_in_l2 = max(l1_oh_start_in_l2 * stride_h - pad_h, 0)
                    l1_iw_start_in_l2 = max(l1_ow_start_in_l2 * stride_w - pad_w, 0)
                    l1_ih_end_in_l2 = min(l1_oh_end_in_l2 * stride_h - pad_h + dilation_h * (kernel_h - 1), in_height - 1)
                    l1_iw_end_in_l2 = min(l1_ow_end_in_l2 * stride_w - pad_w + dilation_w * (kernel_w - 1), in_width - 1)
                    l1_ic_end_in_l2 = min(l1_ic_start_in_l2 + IC_L1, in_channel) - 1
                    ih_size_l1 = l1_ih_end_in_l2 - l1_ih_start_in_l2 + 1
                    iw_size_l1 = l1_iw_end_in_l2 - l1_iw_start_in_l2 + 1
                    ic_size_l1 = l1_ic_end_in_l2 - l1_ic_start_in_l2 + 1
                    iC_size_l1 = ceil_div(ic_size_l1, K0)
                    ifm_CHWc_l1 = HWC2CHWc_DDR2L1(ifm_HWC, l1_ic_start_in_l2, l1_ih_start_in_l2, l1_iw_start_in_l2, ic_size_l1, ih_size_l1, iw_size_l1)
                    k_size_l1 = kernel_h * kernel_w * ceil_div(ic_size_l1, K0) * K0
                    k_start_l1 = kernel_h * kernel_w * l1_ic_start_in_l2
                    weight_MK_l1 = MK2K1MK0_DDR2L1(weight_MK, l1_oc_start_in_l2, k_start_l1, oc_size_l1, k_size_l1)
                    for l0_oc_start_in_l1 in range(0, oc_size_l1, OC_L0):
                        for l0_oh_start_in_l1 in range(0, oh_size_l1, H_L0):
                            for l0_ow_start_in_l1 in range(0, ow_size_l1, W_L0):
                                for l0_iC_start_in_l1 in range(0, iC_size_l1, IC_L0):
                                    l0_oc_start_in_l2 = l1_oc_start_in_l2 + l0_oc_start_in_l1
                                    l0_oh_start_in_l2 = l1_oh_start_in_l2 + l0_oh_start_in_l1
                                    l0_ow_start_in_l2 = l1_ow_start_in_l2 + l0_ow_start_in_l1
                                    l0_oc_end_in_l2 = min(l0_oc_start_in_l2 + OC_L0, l1_oc_start_in_l2 + oc_size_l1) - 1
                                    l0_oh_end_in_l2 = min(l0_oh_start_in_l2 + H_L0, l1_oh_start_in_l2 + oh_size_l1) - 1
                                    l0_ow_end_in_l2 = min(l0_ow_start_in_l2 + W_L0, l1_ow_start_in_l2 + ow_size_l1) - 1
                                    l0_ih_start_in_l2 = max(l0_oh_start_in_l2 * stride_h - pad_h, 0)
                                    l0_iw_start_in_l2 = max(l0_ow_start_in_l2 * stride_w - pad_w, 0)
                                    l0_ih_end_in_l2 = min(l0_oh_end_in_l2 * stride_h - pad_h + dilation_h * (kernel_h - 1), in_height - 1)
                                    l0_iw_end_in_l2 = min(l0_ow_end_in_l2 * stride_w - pad_w + dilation_w * (kernel_w - 1), in_width - 1)
                                    l0_ih_start_in_l1 = l0_ih_start_in_l2 - l1_ih_start_in_l2
                                    l0_ih_end_in_l1 = l0_ih_end_in_l2 - l1_ih_start_in_l2
                                    l0_iw_start_in_l1 = l0_iw_start_in_l2 - l1_iw_start_in_l2
                                    l0_iw_end_in_l1 = l0_iw_end_in_l2 - l1_iw_start_in_l2
                                    l0_iw_size = l0_iw_end_in_l1 - l0_iw_start_in_l1 + 1
                                    l0_ih_size = l0_ih_end_in_l1 - l0_ih_start_in_l1 + 1
                                    l0_ow_size = l0_ow_end_in_l2 - l0_ow_start_in_l2 + 1
                                    l0_oh_size = l0_oh_end_in_l2 - l0_oh_start_in_l2 + 1
                                    l0_iC_size = min(iC_size_l1 - l0_iC_start_in_l1, IC_L0)
                                    l0_oc_size = min(oc_size_l1 - l0_oc_start_in_l1, OC_L0)
                                    if(l0_oh_start_in_l2 * stride_h - pad_h < 0):
                                        l0_pad_t = pad_h - l0_oh_start_in_l2 * stride_h
                                    else:
                                        l0_pad_t = 0
                                    if(l0_oh_end_in_l2 * stride_h - pad_h + dilation_h * (kernel_h - 1) >= in_height):
                                        l0_pad_b = l0_oh_end_in_l2 * stride_h - pad_h + dilation_h * (kernel_h - 1) - in_height + 1
                                    else:
                                        l0_pad_b = 0
                                    if(l0_ow_start_in_l2 * stride_w - pad_w < 0):
                                        l0_pad_l = pad_w - l0_ow_start_in_l2 * stride_w
                                    else:
                                        l0_pad_l = 0
                                    if(l0_ow_end_in_l2 * stride_w - pad_w + dilation_w * (kernel_w - 1) >= in_width):
                                        l0_pad_r = l0_ow_end_in_l2 * stride_w - pad_w + dilation_w * (kernel_w - 1) - in_width + 1
                                    else:
                                        l0_pad_r = 0
                                    bias_en = (l1_ic_start_in_l2==0) and (l0_iC_start_in_l1==0)
                                    psum_en = (l1_ic_start_in_l2 > 0) or (l0_iC_start_in_l1 > 0)
                                    deq_en = (l1_ic_start_in_l2 // IC_L1 == ceil_div(in_channel, IC_L1) - 1) and (l0_iC_start_in_l1 // IC_L0 == ceil_div(iC_size_l1, IC_L0) - 1)
                                    psb_addr = l0_oh_start_in_l1//H_L0 * ceil_div(ow_size_l1, W_L0) * ceil_div(oc_size_l1, OC_L0) + l0_ow_start_in_l1//W_L0 * ceil_div(oc_size_l1, OC_L0) + l0_oc_start_in_l1//OC_L0
                                    ifm_M1K1M0K0 = CHWc2M1K1M0K0_L12L0(ifm_CHWc_l1, l0_ih_size, l0_iw_size, l0_iC_size, l0_ih_start_in_l1, l0_iw_start_in_l1, l0_iC_start_in_l1, kernel_h, kernel_w, stride_h, stride_w, l0_pad_t, l0_pad_b, l0_pad_l, l0_pad_r, dilation_h, dilation_w)
                                    weight_N1K1N0K0_l0 = K1NK02N1K1N0K0_L12RMB(weight_MK_l1, l0_oc_start_in_l1, l0_iC_start_in_l1*kernel_h*kernel_w, l0_oc_size, l0_iC_size*kernel_h*kernel_w)
                                    bias_N1N0 = N2N1N0_L12MMB(bias, l0_oc_start_in_l2, l0_oc_size)
                                    deq_N1N0 = N2N1N0_L12MMB(deq, l0_oc_start_in_l2, l0_oc_size)
                                    result_m2n2m1n1m0n0_psb[psb_addr] = matmul_m1k1m0k0_n1k1n0k0(\
                                    result_m2n2m1n1m0n0_psb[psb_addr], ifm_M1K1M0K0, weight_N1K1N0K0_l0, bias_N1N0, deq_N1N0, ifm_M1K1M0K0.shape[0], weight_N1K1N0K0_l0.shape[0], ifm_M1K1M0K0.shape[1], bias_en, psum_en, deq_en)
                                    if(deq_en):
                                        ofm_MN = CONV_M1N1M0N02MN(result_m2n2m1n1m0n0_psb[psb_addr])
                                        ofm_CHW = MN2CHW(ofm_MN, l0_oc_size, l0_oh_size, l0_ow_size)
                                        ofm[l0_oc_start_in_l2:l0_oc_end_in_l2 + 1, l0_oh_start_in_l2:l0_oh_end_in_l2 + 1, l0_ow_start_in_l2:l0_ow_end_in_l2 + 1] = ofm_CHW
    ofm_golden = Conv2d(ifm, weight, bias, deq, layer)
    diff = ofm_golden.flatten() - ofm.flatten()
    if(np.abs(diff).sum() == 0.0):
        print("img2col + gemm pass")
    else:
        print("img2col + gemm fail")
    return True

test_num = 50
while(test_num > 0):
    if(test_convolution()):
        test_num -= 1